import os
import keras
from keras import backend as K
import tensorflow.compat.v1 as tf
from keras.models import Model
from keras.layers import LSTM,Input, Embedding, Dense, TimeDistributed, Dropout, Bidirectional, Lambda
#from keras_contrib.layers.crf import CRF
from tensorflow.python.tools import freeze_graph
from keras_contrib.layers.crf import CRF, crf_loss, crf_viterbi_accuracy
#K.set_learning_phase(0);
print(keras.name)

def create_model():
    maxlen = 100
    n_tags = 100
    n_words = 100
    input = Input(shape=(maxlen,), name="input")
    model = Embedding(input_dim=n_words + 1, output_dim=250, input_length=maxlen)(input)
    model = Bidirectional(LSTM(units=150, return_sequences=True, recurrent_dropout=0.1))(model) # variational biLSTM
    model = TimeDistributed(Dense(150, activation="relu"))(model)
    crf = CRF(n_tags) # CRF layer
    out = crf(model) # output
    model = Model(input, out)
    model.compile(optimizer="rmsprop", loss=crf.loss_function, metrics=[crf.accuracy])
    model.summary()
    return model

model = create_model()

save_dir = "C:\Users\rahul\PycharmProjects\XYZ\model"

tf.saved_model.simple_save(K.get_session(),
                            save_dir,
                            inputs={"input": model.inputs[0]},
                            outputs={"output": model.outputs[0]}
                            )

freeze_graph.freeze_graph(None, None, None, None, model.outputs[0].op.name, None, None,
                            os.path.join(save_dir, "frozen_model.pb"), False, "", input_saved_model_dir=save_dir)


##########
# python -m tf2onnx.convert --input frozen_model.pb --output model.onnx --outputs crf_1/cond/Merge:0 --inputs input:0 -- opset 11
# 2020-04-30 09:49:37,124 - INFO - Using tensorflow=1.14.0, onnx=1.6.0, tf2onnx=1.6.0/82f805
# 2020-04-30 09:49:37,124 - INFO - Using opset <onnx, 11>
# 2020-04-30 09:49:37,938 - INFO - Optimizing ONNX model
# 2020-04-30 09:49:39,080 - INFO - After optimization: And -1 (6->5), Cast -29 (61->32), Concat -1 (7->6), Const -153 (224->71), Expand -1 (6->5), Identity -27 (36->9), Less -3 (14->11), MatMul -1 (5->4), Mul -1 (12->11), ReduceSum -1 (2->1), Reshape -2 (6->4), Shape -4 (17->13), Slice -1 (28->27), Split -1 (2->1), Squeeze -7 (19->12), Tile -3 (4->1), Transpose -1 (15->14), Unsqueeze -10 (20->10)
# 2020-04-30 09:49:39,111 - INFO -
# 2020-04-30 09:49:39,111 - INFO - Successfully converted TensorFlow model frozen_model.pb to ONNX
# 2020-04-30 09:49:39,111 - INFO - ONNX model is saved at model.onnx